package rtpencoder

import (
	"fmt"

	"gitee.com/general252/gortc/pkg/format/codec"
	"github.com/pion/rtp"
	"github.com/yapingcat/gomedia/go-mpeg2"
)

type _psEncoder struct {
	// payload type of packets.
	PayloadType uint8

	// SSRC of packets (optional).
	// It defaults to a random value.
	SSRC *uint32

	// initial sequence number of packets (optional).
	// It defaults to a random value.
	InitialSequenceNumber *uint16

	// maximum size of packet payloads (optional).
	// It defaults to 1460.
	PayloadMaxSize int

	Codecs []codec.CodecId

	sequenceNumber uint16

	muxer       *mpeg2.PSMuxer
	cidStreamId map[codec.CodecId]uint8
	rtpEncode   *_rtpEncoder
}

// Init initializes the encoder.
func (e *_psEncoder) Init() error {

	var getCodecId = func(id codec.CodecId) mpeg2.PS_STREAM_TYPE {
		cid := mpeg2.PS_STREAM_H264
		switch id {
		case codec.H264:
			cid = mpeg2.PS_STREAM_H264
		case codec.H265:
			cid = mpeg2.PS_STREAM_H265
		case codec.AAC:
			cid = mpeg2.PS_STREAM_AAC
		case codec.G711A:
			cid = mpeg2.PS_STREAM_G711A
		case codec.G711U:
			cid = mpeg2.PS_STREAM_G711U
		}

		return cid
	}

	e.muxer = mpeg2.NewPsMuxer()
	e.cidStreamId = map[codec.CodecId]uint8{}

	for _, id := range e.Codecs {
		streamId := e.muxer.AddStream(getCodecId(id))
		e.cidStreamId[id] = streamId
	}

	e.rtpEncode = &_rtpEncoder{
		PayloadType:           e.PayloadType,
		SSRC:                  e.SSRC,
		InitialSequenceNumber: e.InitialSequenceNumber,
		PayloadMaxSize:        e.PayloadMaxSize,
	}

	return e.rtpEncode.Init()
}

func (e *_psEncoder) Encode(codec codec.CodecId, nalus []byte) ([]*rtp.Packet, error) {

	var packets []*rtp.Packet
	e.muxer.OnPacket = func(pkg []byte) {
		if pkts, err := e.rtpEncode.Encode(pkg); err == nil {
			packets = append(packets, pkts...)
		}
	}

	streamId, ok := e.cidStreamId[codec]
	if !ok {
		return nil, fmt.Errorf("no stream")
	}

	err := e.muxer.Write(streamId, nalus, 0, 0)
	if err != nil {
		return nil, err
	}

	return packets, nil
}

type _psDecoder struct {
	rtpDecode *_rtpDecoder

	demuxer *mpeg2.PSDemuxer
}

// Init initializes the decoder.
func (d *_psDecoder) Init() error {
	d.demuxer = mpeg2.NewPSDemuxer()

	return d.rtpDecode.Init()
}

// Decode decodes an audio frame from a RTP packet.
func (d *_psDecoder) Decode(pkt *rtp.Packet, onFrame func(frame []byte, codec codec.CodecId, pts uint64, dts uint64)) error {
	var getCodecId = func(t mpeg2.PS_STREAM_TYPE) codec.CodecId {
		switch t {
		case mpeg2.PS_STREAM_AAC:
			return codec.AAC
		case mpeg2.PS_STREAM_H264:
			return codec.H264
		case mpeg2.PS_STREAM_H265:
			return codec.H265
		case mpeg2.PS_STREAM_G711A:
			return codec.G711A
		case mpeg2.PS_STREAM_G711U:
			return codec.G711U
		default:
			return codec.File
		}
	}

	data, err := d.rtpDecode.Decode(pkt)
	if err != nil {
		return err
	}

	d.demuxer.OnFrame = func(frame []byte, cid mpeg2.PS_STREAM_TYPE, pts uint64, dts uint64) {
		id := getCodecId(cid)
		onFrame(frame, id, pts, dts)
	}

	err = d.demuxer.Input(data)
	return err
}
